from torch.utils.data import Dataset
import os
from random import shuffle
from PIL import Image
from config import parser
from torchvision import transforms
import numpy as np


dogs=['beagle', 'dachshund', 'dalmatian', 'jindo', 'maltese', 'pomeranian', 'retriever', 'ShihTzu', 'toypoodle', 'Yorkshireterrier']
args=parser.parse_args()

class Dog10Dataset(Dataset):
    def __init__(self,path,data_type=None):
        super(Dog10Dataset, self).__init__()
        self.path=path
        self.data_label=self.collect_image_label(self.path)
        self.data_type=data_type

    def __len__(self):
        return len(self.data_label)

    def __getitem__(self, index):
        if index==0:
            shuffle(self.data_label)
        n=len(self.data_label)
        index=index%n
        img,y=self.get_image_label(self.data_label[index])

        # if args.use_aug and self.data_type=='train':
        #     img=self.img_augment(img)
        img=img.resize((224,224),Image.BICUBIC)
        img=np.array(img,dtype=np.float32)
        img=np.transpose(img/255.0)
        return img,y

    def get_image_label(self,line):
        image_path,label=line.split('|$')
        label=int(label)
        image=Image.open(image_path).convert('RGB')
        return image,label


    def collect_image_label(self,path):
        lines=[]

        for dir in os.listdir(path):
            for image in os.listdir(path+'/'+dir):
               lines.append(path+'/'+dir+'/'+image+'|$'+str(dogs.index(dir)))

        return lines

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def img_augment(self, image):

        # 随机位置裁剪
        random_crop = self.rand() < 0.5
        # 中心裁剪
        center_crop = self.rand() < 0.5
        # 填充后随机裁剪
        random_crop_padding = self.rand() < 0.5
        # 水平翻转
        h_flip = self.rand() < 0.5
        # 竖直翻转
        v_flip = self.rand() < 0.5
        # 亮度
        bright = self.rand() < 0.5
        # 对比度
        contrast = self.rand() < 0.5
        # 饱和度
        saturation = self.rand() < 0.5
        # 颜色随机变换
        color = self.rand() < 0.5
        compose = self.rand() < 0.5
        # 旋转30
        rotate = self.rand() < 0.5

        if h_flip:
            image = transforms.RandomHorizontalFlip()(image)
        if v_flip:
            image = transforms.RandomVerticalFlip()(image)
        if rotate:
            image = transforms.RandomRotation(30)(image)
        if bright:
            image = transforms.ColorJitter(brightness=1)(image)
        if contrast:
            image = transforms.ColorJitter(contrast=1)(image)
        if saturation:
            image = transforms.ColorJitter(saturation=1)(image)
        if color:
            image = transforms.ColorJitter(hue=0.5)(image)
        if compose:
            image = transforms.ColorJitter(0.5, 0.5, 0.5)(image)
        if random_crop:
            image = transforms.RandomCrop(100)(image)
        if center_crop:
            image = transforms.CenterCrop(100)(image)
        if random_crop_padding:
            image = transforms.RandomCrop(100, padding=8)(image)

        return image



